/*
 * Search.cpp
 *
 *  Created on: 16 Nov 2015
 *      Author: hieu
 */
#include <boost/foreach.hpp>
#include "Search.h"
#include "Stack.h"
#include "../Manager.h"
#include "../Hypothesis.h"
#include "../TrellisPath.h"
#include "../Sentence.h"
#include "../../TrellisPaths.h"
#include "../../InputPathsBase.h"
#include "../../InputPathBase.h"
#include "../../System.h"
#include "../../TranslationTask.h"
#include "../../legacy/Util2.h"
#include "../../PhraseBased/TargetPhrases.h"

using namespace std;

namespace Moses2
{

namespace NSCubePruningMiniStack
{

////////////////////////////////////////////////////////////////////////
Search::Search(Manager &mgr) :
        		Moses2::Search(mgr), m_stack(mgr), m_cubeEdgeAlloc(mgr.GetPool())

, m_queue(QueueItemOrderer(),
		std::vector<QueueItem*, MemPoolAllocator<QueueItem*> >(
				MemPoolAllocator<QueueItem*>(mgr.GetPool())))

, m_seenPositions(
		MemPoolAllocator<CubeEdge::SeenPositionItem>(mgr.GetPool()))

, m_queueItemRecycler(MemPoolAllocator<QueueItem*>(mgr.GetPool()))

{
}

Search::~Search()
{
}

void Search::Decode()
{
	const Sentence &sentence = static_cast<const Sentence&>(mgr.GetInput());

	// init cue edges
	m_cubeEdges.resize(sentence.GetSize() + 1);
	for (size_t i = 0; i < m_cubeEdges.size(); ++i) {
		m_cubeEdges[i] = new (mgr.GetPool().Allocate<CubeEdges>()) CubeEdges(
				m_cubeEdgeAlloc);
	}

	const Bitmap &initBitmap = mgr.GetBitmaps().GetInitialBitmap();
	Hypothesis *initHypo = Hypothesis::Create(mgr.GetSystemPool(), mgr);
	initHypo->Init(mgr, mgr.GetInputPaths().GetBlank(), mgr.GetInitPhrase(),
			initBitmap);
	initHypo->EmptyHypothesisState(mgr.GetInput());
	//cerr << "initHypo=" << *initHypo << endl;

	m_stack.Add(initHypo, mgr.GetHypoRecycle(), mgr.arcLists);
	PostDecode(0);

	for (size_t stackInd = 1; stackInd < sentence.GetSize() + 1;
			++stackInd) {
		//cerr << "stackInd=" << stackInd << endl;
		m_stack.Clear();
		Decode(stackInd);
		PostDecode(stackInd);

		//m_stack.DebugCounts();
		//cerr << m_stacks << endl;
	}

}

void Search::Decode(size_t stackInd)
{
	Recycler<HypothesisBase*> &hypoRecycler = mgr.GetHypoRecycle();

	// reuse queue from previous stack. Clear it first
	std::vector<QueueItem*, MemPoolAllocator<QueueItem*> > &container = Container(
			m_queue);
	//cerr << "container=" << container.size() << endl;
	BOOST_FOREACH(QueueItem *item, container){
		// recycle unused hypos from queue
		Hypothesis *hypo = item->hypo;
		hypoRecycler.Recycle(hypo);

		// recycle queue item
		m_queueItemRecycler.push_back(item);
	}
	container.clear();

	m_seenPositions.clear();

	// add top hypo from every edge into queue
	CubeEdges &edges = *m_cubeEdges[stackInd];

	BOOST_FOREACH(CubeEdge *edge, edges){
		//cerr << *edge << " ";
		edge->CreateFirst(mgr, m_queue, m_seenPositions, m_queueItemRecycler);
	}

	/*
 cerr << "edges: ";
 boost::unordered_set<const Bitmap*> uniqueBM;
 BOOST_FOREACH(CubeEdge *edge, edges) {
 uniqueBM.insert(&edge->newBitmap);
 //cerr << *edge << " ";
 }
 cerr << edges.size() << " " << uniqueBM.size();
 cerr << endl;
	 */

	size_t pops = 0;
	while (!m_queue.empty() && pops < mgr.system.options.cube.pop_limit) {
		// get best hypo from queue, add to stack
		//cerr << "queue=" << queue.size() << endl;
		QueueItem *item = m_queue.top();
		m_queue.pop();

		CubeEdge *edge = item->edge;

		// add hypo to stack
		Hypothesis *hypo = item->hypo;

		if (mgr.system.options.cube.lazy_scoring) {
			hypo->EvaluateWhenApplied();
		}

		//cerr << "hypo=" << *hypo << " " << hypo->GetBitmap() << endl;
		m_stack.Add(hypo, hypoRecycler, mgr.arcLists);

		edge->CreateNext(mgr, item, m_queue, m_seenPositions, m_queueItemRecycler);

		++pops;
	}

	// create hypo from every edge. Increase diversity
	if (mgr.system.options.cube.diversity) {
		while (!m_queue.empty()) {
			QueueItem *item = m_queue.top();
			m_queue.pop();

			if (item->hypoIndex == 0 && item->tpIndex == 0) {
				// add hypo to stack
				Hypothesis *hypo = item->hypo;
				//cerr << "hypo=" << *hypo << " " << hypo->GetBitmap() << endl;
				m_stack.Add(hypo, hypoRecycler, mgr.arcLists);
			}
		}
	}
}

void Search::PostDecode(size_t stackInd)
{
	MemPool &pool = mgr.GetPool();

	const InputPaths &paths = mgr.GetInputPaths();
	const Matrix<InputPath*> &pathMatrix = paths.GetMatrix();
	size_t inputSize = pathMatrix.GetRows();
	size_t numPaths = pathMatrix.GetCols();

	BOOST_FOREACH(const Stack::Coll::value_type &val, m_stack.GetColl()){
		const Bitmap &hypoBitmap = *val.first.first;
		size_t firstGap = hypoBitmap.GetFirstGapPos();
		size_t hypoEndPos = val.first.second;
		//cerr << "key=" << hypoBitmap << " " << firstGap << " " << inputSize << endl;

		// create edges to next hypos from existing hypos
		for (size_t startPos = firstGap; startPos < inputSize; ++startPos) {
			for (size_t pathInd = 0; pathInd < numPaths; ++pathInd) {
				const InputPath *path = pathMatrix.GetValue(startPos, pathInd);

				if (path == NULL) {
					break;
				}
				if (path->GetNumRules() == 0) {
					continue;
				}

				const Range &pathRange = path->range;
				//cerr << "pathRange=" << pathRange << endl;
				if (!CanExtend(hypoBitmap, hypoEndPos, pathRange)) {
					continue;
				}

				const ReorderingConstraint &reorderingConstraint = mgr.GetInput().GetReorderingConstraint();
				if (!reorderingConstraint.Check(hypoBitmap, startPos, pathRange.GetEndPos())) {
					continue;
				}

				const Bitmap &newBitmap = mgr.GetBitmaps().GetBitmap(hypoBitmap, pathRange);
				size_t numWords = newBitmap.GetNumWordsCovered();

				CubeEdges &edges = *m_cubeEdges[numWords];

				// sort hypo for a particular bitmap and hypoEndPos
				const Hypotheses &sortedHypos = val.second->GetSortedAndPruneHypos(mgr, mgr.arcLists);

				size_t numPt = mgr.system.mappings.size();
				for (size_t i = 0; i < numPt; ++i) {
					const TargetPhrases *tps = path->targetPhrases[i];
					if (tps && tps->GetSize()) {
						CubeEdge *edge = new (pool.Allocate<CubeEdge>()) CubeEdge(mgr, sortedHypos, *path, *tps, newBitmap);
						edges.push_back(edge);
					}
				}
			}
		}
	}
}

const Hypothesis *Search::GetBestHypo() const
{
	const Hypothesis *bestHypo = m_stack.GetBestHypo();
	return bestHypo;
}

void Search::AddInitialTrellisPaths(TrellisPaths<TrellisPath> &paths) const
{
	const Stack::Coll &coll = m_stack.GetColl();
	BOOST_FOREACH(const Stack::Coll::value_type &val, coll){
		const Moses2::HypothesisColl &hypos = *val.second;
		const Hypotheses &sortedHypos = hypos.GetSortedAndPruneHypos(mgr, mgr.arcLists);

		BOOST_FOREACH(const HypothesisBase *hypoBase, sortedHypos) {
			const Hypothesis *hypo = static_cast<const Hypothesis*>(hypoBase);
			TrellisPath *path = new TrellisPath(hypo, mgr.arcLists);
			paths.Add(path);
		}
	}
}

}

}

